import os

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image

from config import Input_shape
from k_means import kmeans
from utils.yolo_utils import letterbox_image

INPUT_SIZE = (Input_shape, Input_shape)


def changePath(root_path, file_name, new_name, str_tgt,
               str_org='D:/1aitest'):
    fnew = open(root_path + new_name, 'w')
    
    with open(root_path + file_name, 'r')as f:
        content = f.readlines()
        for index, line in enumerate(content):
            new_line = line.replace(str_org, str_tgt, 1)
            fnew.write(new_line)
    
    fnew.close()


def show_data(root_path, index_file, input_shape, save_path=None, anchor_num=None):
    """
    processes the data into standard shape
    :param annotation_path: path_to_image box1,box2,...,boxN with boxX: x_min,y_min,x_max,y_max,class_index
    :param data_path: saver at "/home/minh/stage/train.npz"
    :param input_shape: (416, 416)
    :param max_boxes: 100: maximum number objects of an image
    :param load_previous: for 2nd, 3th, .. using
    :return: image_data [N, 416, 416, 3] not yet normalized, N: number of image
             box_data: box format: [N, 100, 6], 100: maximum number of an image
                                                6: top_left{x_min,y_min},bottom_right{x_max,y_max},class_index (no space)
                                                /home/minh/keras-yolo3/VOCdevkit/VOC2007/JPEGImages/000012.jpg 156,97,351,270,6
    """
    annotation_path = os.path.join(root_path, index_file)
    with open(annotation_path) as f:
        GG = f.readlines()
        # np.random.shuffle(GG)
        box_kmean = []
        
        for num, line in enumerate(GG):
            if line == '\n':
                continue
            if line.startswith('#'):
                continue
            
            # if num>0:
            #     break
            line = line.split('.png')
            filepath = line[0] + '.png'
            image = Image.open(filepath)
            name_list = filepath.split('/')[-3:]
            filename = '/'.join(name_list)
            # print(num,'  ',filename)
            # For the case 2
            # boxed_image is the resized image of (416, 416), shape_image is original size
            boxed_image, shape_image = letterbox_image(image, tuple(reversed(input_shape)))
            fig, ax = plt.subplots(1)
            ax.imshow(boxed_image)
            
            ax.annotate(filename, (0, 10), color='b',
                        weight='bold',
                        fontsize=9, ha='left', va='center')
            
            info_line = line[1]
            info_line = info_line.split()
            
            for i, box in enumerate(info_line):
                left, top, right, bottom, seedx, seedy, label = list(map(int, box.split(',')))
                
                image_size = np.array(image.size)
                input_size = np.array(input_shape[::-1])  # inverse 没意义
                # for case 2
                # image size after crop w,h
                new_size = (image_size * np.min(input_size / image_size)).astype(np.int32)
                # Correct BB to new image
                left_top = (np.array([left, top]) * new_size / image_size + (input_size - new_size) / 2).astype(
                    np.int32)
                right_bottom = (np.array([right, bottom]) * new_size / image_size + (input_size - new_size) / 2).astype(
                    np.int32)
                seed = (np.array([seedx, seedy]) * new_size / image_size + (input_size - new_size) / 2).astype(np.int32)
                # https://stackoverflow.com/questions/37435369/matplotlib-how-to-draw-a-rectangle-on-image
                # 水平 垂直
                width = right_bottom[0] - left_top[0]
                height = right_bottom[1] - left_top[1]
                if anchor_num is not None:
                    box_kmean.append(np.array([width, height]))
                else:
                    rect = patches.Rectangle(left_top, width, height, linewidth=1, edgecolor='r', facecolor='none')
                    ax.add_patch(rect)
                    label_str = '%d' % label
                    ax.annotate(label_str, (left_top[0] + width / 5, left_top[1] + height * 4 / 5),
                                weight='bold',
                                fontsize=8, ha='center', va='center')
                    ax.annotate(str(i), (left_top[0] + width / 5, left_top[1] + height / 5),
                                weight='bold',
                                fontsize=8, ha='center', va='center')
                    
                    plt.plot([seed[0]], [seed[1]], "ro")
            if anchor_num is None:
                if save_path is not None:
                    save_name = os.path.join(save_path, name_list[-3] + '%s.png' % num)
                    plt.savefig(save_name, format='png')
                else:
                    plt.show()
                    plt.close()
        if anchor_num is not None:
            box_kmean = np.stack(box_kmean, axis=0)
            clusters = kmeans(box_kmean, anchor_num)
            areas = clusters[:, 0] * clusters[:, 1]
            ind = np.argsort(areas)
            clusters = clusters[ind]
            print(clusters)
            list_data = list(clusters.reshape([-1]))
            output = str([int(i) for i in list_data])
            output = output.lstrip('[')
            output = output.rstrip(']')
            with open(root_path + '/new_anchor.txt', 'w')as f:
                f.write(output)


def check_Change_result(root_path, index_file):
    annotation_path = root_path + index_file
    with open(annotation_path) as f:
        GG = f.readlines()
        for i, line in enumerate(GG):
            if line == '\n':
                continue
            if line.startswith('#'):
                continue
            
            num_1 = int((GG[i].split('.png')[1]).split(',')[0])
            if num_1 > 300:
                print(i + 1)


def changePointOrder(root_path, index_file, new_name):
    '''
    一定要在变换路径前操作，将变换完的文件变换路径
    :param root_path:
    :param index_file:
    :param new_name:
    :return:
    '''
    annotation_path = root_path + index_file
    
    fnew = open(root_path + new_name, 'w')
    
    def need_switch(filename):
        image_name = filename.split('/')[-1]
        mirror_path = filename.split('/')[-2]
        is_mirror = False
        is_up = False
        if image_name.startswith('up'):
            is_up = True
        if mirror_path.endswith('mirro'):
            is_mirror = True
        return is_up, is_mirror
    
    def switchY(box_str, is_up, is_mirror):
        x1, y1, x2, y2, seedx, seedy, label = box_str.split(',')
        if is_up and not is_mirror:
            new_box_str = x1 + ',' + y2 + ',' + x2 + ',' + y1 + ',' + seedx + ',' + seedy + ',' + label
        
        elif not is_up and is_mirror:
            new_box_str = x2 + ',' + y1 + ',' + x1 + ',' + y2 + ',' + seedx + ',' + seedy + ',' + label
        
        elif is_up and is_mirror:
            new_box_str = x2 + ',' + y2 + ',' + x1 + ',' + y1 + ',' + seedx + ',' + seedy + ',' + label
        
        else:
            new_box_str = box_str
        
        return new_box_str
    
    with open(annotation_path)as f:
        content = f.readlines()
        # new_line=None
        for index, ori_line in enumerate(content):
            line = ori_line.split('.png')
            filename = line[0] + '.png'
            is_up, is_mirror = need_switch(filename)
            info_line = line[1]
            info_line = info_line.split()
            if is_up or is_mirror:
                boxes_str = ''
                for box in info_line:
                    new_box = switchY(box, is_up, is_mirror)
                    boxes_str += (' ' + new_box)
                new_line = filename + boxes_str + '\n'
            
            
            else:
                new_line = ori_line
            fnew.write(new_line)


def change_order(root_path, index_file, new_name):
    annotation_path = root_path + index_file
    
    fnew = open(root_path + new_name, 'w')
    
    def need_switch(filename):
        image_name = filename.split('/')[-1]
        mirror_path = filename.split('/')[-2]
        is_mirror = False
        is_up = False
        if image_name.startswith('up'):
            is_up = True
        if mirror_path.endswith('mirro'):
            is_mirror = True
        return is_up, is_mirror
    
    def assign_label(box_str, new_label):
        elements_list = box_str.split(',')
        elements_list[-1] = str(new_label)
        return ','.join(elements_list)
    
    with open(annotation_path)as f:
        content = f.readlines()
        # new_line=None
        for index, ori_line in enumerate(content):
            line = ori_line.split('.png')
            filename = line[0] + '.png'
            is_up, is_mirror = need_switch(filename)
            info_line = line[1]
            box_list = info_line.split()
            label_list = []
            if is_mirror:
                for box in box_list:
                    elements_list = box.split(',')
                    label = int(elements_list[-1])
                    if is_up:
                        new_label = 17 - label
                    else:
                        new_label = 49 - label
                    label_list.append(new_label)
                label_list.reverse()
                boxes_str = ''
                for box, new_label in zip(box_list, label_list):
                    new_box = assign_label(box, new_label)
                    boxes_str += (' ' + new_box)
                new_line = filename + boxes_str + '\n'
            else:
                new_line = ori_line
            fnew.write(new_line)


def changeMirrorOrder(root_path, index_file, new_name):
    annotation_path = root_path + index_file
    
    fnew = open(root_path + new_name, 'w')
    
    with open(annotation_path) as f:
        GG = f.readlines()
        line_num = len(GG)
        
        def write_line(f, i, switch):
            ori_line = GG[i]
            if switch:
                line = ori_line.split('.png')
                filename = line[0] + '.png'
                info_line = line[1]
                info_line = info_line.split()
                info_line.reverse()
                reverse_str = ''
                for box in info_line:
                    reverse_str += ' '
                    reverse_str += box
                new_line = filename + reverse_str + '\n'
            else:
                new_line = ori_line
            f.write(new_line)
        
        for i in range(0, line_num, 2):
            
            # info_line = GG[i].split('.png')[1]
            # info_line = info_line.split()
            # num_1=info_line[0].split(',')[0]
            # num_2=info_line[-1].split(',')[0]
            
            num_1 = int((GG[i].split('.png')[1]).split(',')[0])
            num_2 = int((GG[i + 1].split('.png')[1]).split(',')[0])
            # 进行以下比较的预设是 镜像和原版里肯定有且只有一组是反的，即从右到左
            if num_1 > num_2:
                switch_1 = True
                switch_2 = False
            else:
                switch_2 = True
                switch_1 = False
            write_line(fnew, i, switch_1)
            write_line(fnew, i + 1, switch_2)
    fnew.close()


def exclude_cases(f_include, f_compare, f_exclude_path):
    '''
    delete data entries present in f_compare from f_include
    :param root_path:
    :param f_include:
    :param f_compare:
    :return:
    '''
    
    fnew = open(f_exclude_path, 'w')
    f_compare = open(f_compare, 'r')
    content_com = f_compare.readlines()
    dup_num = 0
    with open(f_include, 'r')as f:
        for index, line in enumerate(f.readlines()):
            if line == '\n':
                continue
            if line.startswith('#'):
                continue
            
            filename = (line.split('.png'))[0].split('/')[-3]
            need_write = True
            
            for i, line_c in enumerate(content_com):
                if filename in line_c:
                    dup_num += 1
                    print(i, '  ', filename)
                    need_write = False
                    break
            if need_write:
                fnew.write(line)
    print('dup_num = %d' % dup_num)
    fnew.close()
    f_compare.close()


def listDir(rootDir, record_name):
    record_name = os.path.join(rootDir, record_name)
    with open(record_name, 'w')as f:
        for filename in os.listdir(rootDir):
            pathname = os.path.join(rootDir, filename)
            f.write(pathname + '\n')


def copyFiles(sourceDir, targetDir):
    if not os.path.exists(targetDir):
        os.makedirs(targetDir)
    
    for i, file in enumerate(os.listdir(sourceDir)):
        sourceFile1 = os.path.join(sourceDir, file, 'Setting Data2/aitest/low.png')
        sourceFile2 = os.path.join(sourceDir, file, 'Setting Data2/aitest/up.png')
        if os.path.exists(sourceFile1):
            targetFile1 = os.path.join(targetDir, 'low%d.png'%(i))
            if not os.path.exists(targetFile1) or (
                    os.path.exists(targetFile1) and (os.path.getsize(targetFile1) != os.path.getsize(sourceFile1))):
                open(targetFile1, "wb").write(open(sourceFile1, "rb").read())
        
        if os.path.isfile(sourceFile2):
            targetFile2 = os.path.join(targetDir, 'up%d.png'%(i))
            if not os.path.exists(targetFile2) or (
                    os.path.exists(targetFile2) and (os.path.getsize(targetFile2) != os.path.getsize(sourceFile2))):
                open(targetFile2, "wb").write(open(sourceFile2, "rb").read())


if __name__ == '__main__':
    # copyFiles('F:/ProjectData/yolo3/test/testpic/AI', 'F:/ProjectData/yolo3/test/testpic/zhangyi')
    
    # listDir('F:/ProjectData/yolo3/test/image','rcd.txt')
    
    # changeMirrorOrder('F:/ProjectData/yolo3/classify/1aitest','/train.txt','/mirror.txt')
    # changePointOrder('F:/ProjectData/yolo3/classify/1aitest', '/mirror.txt', '/point.txt')
    # check_Change_result('F:/ProjectData/yolo3/1aitest','/train_new.txt')
    
    # include_path='F:/ProjectData/yolo3/newData/AIData3/1aitest-0621/ai pixel value.txt'
    # compare_path='F:/ProjectData/yolo3/newData/AIData3/1aitest-0621/ai pixel value2.txt'
    # new_path='F:/ProjectData/yolo3/newData/AIData3/1aitest-0621/add.txt'
    # exclude_cases(f_include=include_path, f_compare=compare_path, f_exclude_path=new_path)
    
    # change_order('F:/ProjectData/yolo3/newdata/oldTool','/ai pixel value.txt','/right_order.txt')
    
    data_path = 'F:/ProjectData/yolo3/newData/AIData3/1aitest'
    #
    # str_org = 'F:/ProjectData/yolo3/newdata/AITest-0420/1aitest'
    # str_tgt = 'F:/ProjectData/yolo3/newdata/AITest_0420/1aitest'
    changePath(root_path=data_path,
               file_name='/ai pixel value.txt',
               new_name='/train.txt',
               # str_org=str_org,
               str_tgt=data_path
               )
    #
    # save_path='F:/ProjectData/yolo3/newdata/AITest-0420/1aitest'
    # show_data(data_path, 'valid_result.txt', INPUT_SIZE, save_path=None, anchor_num=None)
